import cv2
from time import time
import copy
from .rw_ini import IniData
from .post_process_distance import *
from .utils_weice import *
from .pmi_bajiao import *
from .pmi_pad2bajiao_distance import *
def post_all_process_aoi(can_wuran, wuran,post_results, ini_data, image,info_dict,per_return_dict, distance, set_num_zhenhen,set_diagonal_length):
    pad = []
    pad_model_list = []
    zhenhen = []
    cutting_path = []
    
    can_wuran_num = []
    zhenhen_crop = []
    madian_pad_crop = []
    
    return_can_wuran = []
    return_zhenhen=[]    
    return_wuran=[]
    return_pad=[]

    # print("post_results: ",post_results)
    
    for i in range(len(post_results)):
        label, score ,box = post_results[i]
        if label == 1:
            pad.append(post_results[i])
        elif label == 2:
            zhenhen.append(post_results[i])
        elif label == 4:
            cutting_path.append(post_results[i])
    


    # can_wuran = keyichu_back_result
    # wuran = wuran_back_result   
    
    # 移除切割道的所有box
    # start_cutting_time = time.time()
    for cur_cutting_path in cutting_path:
        cur_cutting_path_box = cur_cutting_path[2]
        cur_cutting_path_box,_ = check_bbox_bounds(image, cur_cutting_path_box)
        for i in range(len(pad)-1,-1,-1):
            cur_pad_box = pad[i][2]
            cur_pad_box,_ = check_bbox_bounds(image, cur_pad_box)
            if is_in_bbox(cur_cutting_path_box,cur_pad_box):
                pad.pop(i)
        for i in range(len(zhenhen)-1,-1,-1):
            cur_zhenhen_box = zhenhen[i][2]
            cur_zhenhen_box,_ = check_bbox_bounds(image, cur_zhenhen_box)
            if is_in_bbox(cur_cutting_path_box,cur_zhenhen_box):
                zhenhen.pop(i)
        for i in range(len(can_wuran)-1,-1,-1):
            cur_can_wuran_box = can_wuran[i][2]
            cur_can_wuran_box,_ = check_bbox_bounds(image, cur_can_wuran_box)
            inter, union_area, area_0, area_1 = iou(cur_cutting_path_box, cur_can_wuran_box)
            if inter/area_1 ==1:
                can_wuran.pop(i)
        for i in range(len(wuran)-1,-1,-1):
            cur_wuran_box = wuran[i][2]
            cur_wuran_box,_ = check_bbox_bounds(image, cur_wuran_box)
            inter, union_area, area_0, area_1 = iou(cur_cutting_path_box, cur_wuran_box)
            if inter/area_1 ==1:
                wuran.pop(i)
    # print(len(can_wuran))
    # print(len(wuran))
    # end_cutting_time = time.time()
    # print("cutting_time: ",end_cutting_time-start_cutting_time)
    # print("wuran: ",wuran)
    # print("can_wuran: ",can_wuran)
    #过滤污染和可移除异物重合的框
    for i in range(len(wuran)-1,-1,-1):
        cur_wuran = wuran[i]
        label_wuran, score_wuran ,box_wuran = cur_wuran
        for j in range(len(can_wuran)-1,-1,-1):
            cur_can_wuran = can_wuran[j]
            label_can_wuran, score_can_wuran ,box_can_wuran = cur_can_wuran
            inter, union_area, area_wuran, area_can_wuran = iou(box_wuran,box_can_wuran)
            if inter/union_area > 0.5:
                if score_wuran > score_can_wuran:
                    can_wuran.pop(j)
                else:
                    wuran.pop(i)

    #关于污染和可移除异物的操作
    # strat_wuran_time = time.time()
    for cur_wuran in wuran:
        label_wuran, score_wuran ,box_wuran = cur_wuran
        box_wuran,_ = check_bbox_bounds(image, box_wuran)
        # 只要是污染就返回box
        FaultX = float(ini_data.X + ini_data.PixelSizeX * ((box_wuran[0] + box_wuran[2]) / 2 - image.shape[1] / 2))
        FaultY = float(ini_data.Y + ini_data.PixelSizeY * ((box_wuran[1] + box_wuran[3]) / 2 - image.shape[0] / 2))
        return_wuran.append([box_wuran,[FaultX, FaultY]])
        
    for cur_can_wuran in can_wuran:
        label_can_wuran, score_can_wuran ,box_can_wuran = cur_can_wuran
        box_can_wuran, _ = check_bbox_bounds(image, box_can_wuran)
        #可移除异物如果大于 set_diagonal_length就计算数量
        diagonal_length = calculate_diagonal_length(box_can_wuran, ini_data.PixelSizeX, ini_data.PixelSizeY)
        # print("diagonal_length: ",diagonal_length)
        if diagonal_length > set_diagonal_length:
            FaultX = float(ini_data.X + ini_data.PixelSizeX * ((box_can_wuran[0] + box_can_wuran[2]) / 2 - image.shape[1] / 2))
            FaultY = float(ini_data.Y + ini_data.PixelSizeY * ((box_can_wuran[1] + box_can_wuran[3]) / 2 - image.shape[0] / 2))
            return_can_wuran.append([box_can_wuran,[FaultX, FaultY]])    
        else:
            FaultX = float(ini_data.X + ini_data.PixelSizeX * ((box_can_wuran[0] + box_can_wuran[2]) / 2 - image.shape[1] / 2))
            FaultY = float(ini_data.Y + ini_data.PixelSizeY * ((box_can_wuran[1] + box_can_wuran[3]) / 2 - image.shape[0] / 2))
            can_wuran_num.append([box_can_wuran,[FaultX, FaultY]])
    if len(can_wuran_num) >3:
        for can_remove in can_wuran_num:
            return_can_wuran.append(can_remove)
           
        
    #过滤图片周围set_pad_dis像素内的pad
    # start_filter_time = time.time()
    set_pad_dis = 5
    image_filter = [set_pad_dis,set_pad_dis,image.shape[1]-set_pad_dis,image.shape[0]-set_pad_dis]
    pad_filter_list = []
    for cur_pad in pad:
        cur_pad_box = cur_pad[2]
        cur_pad_box,_ = check_bbox_bounds(image, cur_pad_box) 
        inter, union_area, area_0, area_1 = iou(image_filter, cur_pad_box)
        if inter/area_1 != 1:
            continue 
        pad_filter_list.append(cur_pad)
    
    
    pad_bbox_list = []
    pad_bbox_new_col_list = []
    pad_bbox_new_list = []
    for cur_pad in pad_filter_list:
        cur_pad_box = cur_pad[2]
        cur_pad_box,_ = check_bbox_bounds(image, cur_pad_box)
        pad_bbox_list.append(cur_pad_box)
    
    ## 更新对齐pad box
    # 将输出的pad list根据列排序
    reslut_col_lists = split_col_pad_list(pad_bbox_list)
    # 计算方差更新x0，x1坐标
    update_pad_col_list = cal_col_var(reslut_col_lists)
    for update_pad_col in update_pad_col_list:
        for update_pad in update_pad_col:
            pad_bbox_new_col_list.append(update_pad)
    # 将修正后的list根据行排序
    reslut_row_lists = split_row_pad_list(pad_bbox_new_col_list)
    # 计算方差更新y0,y1坐标 
    update_pad_row_list = cal_row_var(reslut_row_lists)
    for update_pad_row in update_pad_row_list:
        for update_pad in update_pad_row:
            pad_bbox_new_list.append(update_pad)

    # image_pad_draw = image.copy()
    # if FLAG_DRAW_PAD_LINE:
    #     for pad_box in pad_bbox_list:
    #         cv2.rectangle(image_pad_draw, (int(pad_box[0]), int(pad_box[1])), (int(pad_box[2]), int(pad_box[3])), (255, 0, 0), 1)
    
    # if FLAG_DRAW_PAD_LINE:
    #     for pad_bbox_new_row in pad_bbox_new_list:
    #         cv2.rectangle(image_pad_draw, (int(pad_bbox_new_row[0]), int(pad_bbox_new_row[1])), (int(pad_bbox_new_row[2]), int(pad_bbox_new_row[3])), (0, 0, 255), 1)
    
    # cv2.imwrite("./fix_images/"+ini_data.filename,image_pad_draw)
    
    
    # start_judge_distance_time = time.time()
    # 根据det出来的bbox计算针痕异常
    for index_pad , cur_pad_box in enumerate(pad_bbox_new_list):
        cur_pad_box,_ = check_bbox_bounds(image, cur_pad_box)
        #更新pad box
        pad_box_new= add_pad_edges(image, cur_pad_box)
        # pad_box_new = cur_pad_box
        
        #将所有的pad box保存用作cls
        madian_pad_crop.append(cur_pad_box)
        # cv2.rectangle(image, (int(box_zhenhen[0]), int(box_zhenhen[1])), (int(box_zhenhen[2]), int(box_zhenhen[3])), (0, 255, 0), 1)
        # cv2.rectangle(image, (int(cur_pad_box[0]), int(cur_pad_box[1])), (int(cur_pad_box[2]), int(cur_pad_box[3])), (255, 0, 0), 1)
        for cur_zhenhen in zhenhen:
            box_zhenhen = cur_zhenhen[2]
            box_zhenhen,_ = check_bbox_bounds(image, box_zhenhen)
            # print(type(box_zhenhen[0]))
            x1, y1, x2, y2 = int(box_zhenhen[0]), int(box_zhenhen[1]), int(box_zhenhen[2]), int(box_zhenhen[3])
            x3, y3, x4, y4 = int(pad_box_new[0]), int(pad_box_new[1]), int(pad_box_new[2]), int(pad_box_new[3])
            #判断针痕是否在pad内部
            if x1 > x3 and y1 > y3 and x2 < x4 and y2 < y4:  
                #计算dist
                dist_x = min(x1 - x3, x4 - x2)*ini_data.PixelSizeX
                dist_y = min(y1 - y3, y4 - y2)*ini_data.PixelSizeY
                dist = min(dist_x, dist_y)
                # print(dist)
                if dist <= distance:  #针痕距离pad边框小于设定值    
                    FaultX = float(ini_data.X + ini_data.PixelSizeX * ((box_zhenhen[0] + box_zhenhen[2]) / 2 - image.shape[1] / 2))
                    FaultY = float(ini_data.Y + ini_data.PixelSizeY * ((box_zhenhen[1] + box_zhenhen[3]) / 2 - image.shape[0] / 2))
                    return_zhenhen.append([box_zhenhen,[FaultX, FaultY],dist,pad_box_new])
                else:
                    # zhenhen_crop.append(pad_bbox_new)  
                    zhenhen_crop.append(pad_box_new) 
            elif boxes_intersection(box_zhenhen, pad_box_new):
                dist=0
                FaultX = float(ini_data.X + ini_data.PixelSizeX * ((box_zhenhen[0] + box_zhenhen[2]) / 2 - image.shape[1] / 2))
                FaultY = float(ini_data.Y + ini_data.PixelSizeY * ((box_zhenhen[1] + box_zhenhen[3]) / 2 - image.shape[0] / 2))
                return_zhenhen.append([box_zhenhen,[FaultX, FaultY],dist,pad_box_new])
        
    
    info_dict['zhenhen_crop'] = zhenhen_crop
    info_dict['madian_pad'] = madian_pad_crop    
    
    per_return_dict["wuran_box"] = return_wuran
    per_return_dict["can_wuran_box"] = return_can_wuran
    per_return_dict["zhenhen_box"] = return_zhenhen
    # print("return_zhenhen: ",len(return_zhenhen)) 
    # per_return_dict["madian_box"] = return_madian_list
    # per_return_dict["zhenhen_num_error"] = return_pad
    return info_dict , per_return_dict
   
def post_all_process_pmi(post_results, image,info_dict,per_return_dict, distance):
    
    bajiao_pad = []
    pad = []
    zhenhen = []
    
    zhenhen_crop=[]

    return_zhenhen=[]    
    return_wuran=[]
    return_pad=[]
    
    for i in range(len(post_results)):
        label, score ,box = post_results[i]
        if label == 2:
            bajiao_pad.append(post_results[i])
        elif label == 0 or label == 1:
            pad.append(post_results[i])
        elif label == 3:
            zhenhen.append(post_results[i])
    
    #过滤图片周围set_pad_dis像素内的pad
    set_pad_dis = 5
    image_filter = [set_pad_dis,set_pad_dis,image.shape[1]-set_pad_dis,image.shape[0]-set_pad_dis]
    pad_filter_list = []
    bajiao_box_list = []
    
    for cur_pad in pad:
        cur_pad_box = cur_pad[2]
        cur_pad_box,_ = check_bbox_bounds(image, cur_pad_box) 
        inter, union_area, area_0, area_1 = iou(image_filter, cur_pad_box)
        if inter/area_1 != 1:
            continue 
        pad_filter_list.append(cur_pad)
   
    #处理常规的pad
    for index_pad , cur_pad in enumerate(pad_filter_list):
        cur_pad_box = cur_pad[2]
        cur_pad_box,_ = check_bbox_bounds(image, cur_pad_box)
        for cur_zhenhen in zhenhen:
            box_zhenhen = cur_zhenhen[2]
            box_zhenhen,_ = check_bbox_bounds(image, box_zhenhen)
            x1, y1, x2, y2 = int(box_zhenhen[0]), int(box_zhenhen[1]), int(box_zhenhen[2]), int(box_zhenhen[3])
            x3, y3, x4, y4 = int(cur_pad_box[0]), int(cur_pad_box[1]), int(cur_pad_box[2]), int(cur_pad_box[3])
            #判断针痕是否在pad内部
            if x1 > x3 and y1 > y3 and x2 < x4 and y2 < y4:  
                #计算dist
                dist_x = min(x1 - x3, x4 - x2)
                dist_y = min(y1 - y3, y4 - y2)
                dist = min(dist_x, dist_y)
                if dist <= distance:  #针痕距离pad边框小于设定值    
                    return_zhenhen.append([box_zhenhen,dist,cur_pad_box])
                else:
                    pad_area = (x4-x3)*(y4-y3)
                    zhenhen_crop.append([cur_pad_box,pad_area]) 
            elif boxes_intersection(box_zhenhen, cur_pad_box):
                dist=0
                return_zhenhen.append([box_zhenhen,dist,cur_pad_box])
                
            
    for cur_bajiao_pad in bajiao_pad:
        cur_bajiao_pad_box = cur_bajiao_pad[2]
        cur_bajiao_pad_box,_ = check_bbox_bounds(image, cur_bajiao_pad_box)
        inter, union_area, area_0, area_1 = iou(image_filter, cur_bajiao_pad_box)
        if inter/area_1 != 1:
            continue
        bajiao_box_list.append(cur_bajiao_pad_box)

    #处理八角形的pad
    result_bajiao_list = []
    if len(bajiao_box_list) > 0:
        result_bajiao_list = get_bajiao(image,bajiao_box_list)
        # 计算针痕和八角的距离
        for result_bajiao,bajiao_box in zip(result_bajiao_list,bajiao_box_list):
            pad_area, bajiao_loc = result_bajiao[0],result_bajiao[1]
            zhenhen_crop.append([bajiao_box,pad_area])
            for cur_zhenhen in zhenhen:
                box_zhenhen = cur_zhenhen[2]
                box_zhenhen,_ = check_bbox_bounds(image, box_zhenhen)
                x1, y1, x2, y2 = int(box_zhenhen[0]), int(box_zhenhen[1]), int(box_zhenhen[2]), int(box_zhenhen[3])
                x3, y3, x4, y4 = int(bajiao_box[0]), int(bajiao_box[1]), int(bajiao_box[2]), int(bajiao_box[3])
                #判断针痕是否在pad内部
                if x1 > x3 and y1 > y3 and x2 < x4 and y2 < y4:
                    #计算dist
                    dist = cal_min_distance(bajiao_loc,box_zhenhen)
                    if dist <= distance:  #针痕距离pad边框小于设定值
                        return_zhenhen.append([box_zhenhen,dist,bajiao_box])
                elif boxes_intersection(box_zhenhen, bajiao_box):
                    dist=0
                    return_zhenhen.append([box_zhenhen,dist,bajiao_box])
    else:
        pass
    
    info_dict['zhenhen_crop'] = zhenhen_crop
        
    per_return_dict["wuran_box"] = return_wuran
    per_return_dict["zhenhen_box"] = return_zhenhen 
    # per_return_dict["zhenhen_num_error"] = return_pad
    return info_dict , per_return_dict     
